/**
 * Copyright 2024/3/23 the original brandnewdata.com.
 * <p>
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * <p>
 * https://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.zlb.envmodify.transform.transformer;

import com.zlb.envmodify.entity.MethodOverWriteConfig;
import com.zlb.envmodify.service.EnvModifyServiceFactory;
import com.zlb.envmodify.transform.config.ModifyClassConfig;
import javassist.*;

import java.io.ByteArrayInputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * @Description 对加载的类进行动态增强
 * @ClassName ClassTransformer
 * @Author zhulb
 * @Date 2024/3/23 0:48
 * @Version 0.0.1
 */
public class ClassTransformer implements ClassFileTransformer {

    static Logger logger = Logger.getLogger(ClassTransformer.class.getName());
    private ModifyClassConfig modifyClassConfig;

    public ClassTransformer() {
        modifyClassConfig = EnvModifyServiceFactory.getSigleton().getEnvModifyClassConfig();
    }

    static Map<ClassLoader, Object> loaders = new HashMap<>();

    @Override
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        try {
//            if (loader == null) {
//                System.out.println(className + " ->系统类加载器");
//            } else {
//                System.out.println(className + " ->" + loader.toString());
//            }
            //1.不需要代理的情况
            if (className == null) {
                return classfileBuffer;
            }
            if (loader != null && loaders.get(loader) == null) {
                if ("sun.misc.Launcher$AppClassLoader".equals(loader.getClass().getName()) || "sun.misc.Launcher.ExtClassLoader".equals(loader.getClass().getName())) {

                } else {
                    //很牛逼的一个代码 问chatgpt->javaassit  传入的classloader有什么用，并举例说明
                    //问题描述：因为SpringBoot使用的自定义的类加载器，
                    //所以对类进行动态修改时，涉及到动态编译，所以会去解析依赖的类，而依赖的类不在ClassPool的
                    //class path路径下会导致报找不到类的错误
                    ClassPool classPool = ClassPool.getDefault();
                    classPool.appendClassPath(new LoaderClassPath(loader));
                }
                loaders.put(loader, 1);
            }
            Map<String, List<MethodOverWriteConfig>> overWriteConfig = modifyClassConfig.getMethodOverWriteConfigs().get(className.replace("/", "."));
            if (overWriteConfig == null) {
                return classfileBuffer;
            }
            //2.需要进行代理
            return doTransform(className, overWriteConfig, classfileBuffer);
        } catch (Exception e) {
            e.printStackTrace();
            System.exit(99);
        }
        return null;
    }

    private byte[] doTransform(String className, Map<String, List<MethodOverWriteConfig>> overWriteConfig, byte[] classfileBuffer) {
        try {
            ClassPool classPool = ClassPool.getDefault();
            // ifNotFrozen参数传false，那么就可以代理几乎所有的类(处理JDK核心的一些类)
            // 比如java.util.Properties的Frozen=true，只有这种方式才能代理成功
            CtClass cl = classPool.makeClass(new ByteArrayInputStream(classfileBuffer), false);
            for (String key : overWriteConfig.keySet()) {
                doTransform(key, className, cl, overWriteConfig);
            }
            return cl.toBytecode();
        } catch (Exception e) {
            logger.log(Level.WARNING, e, () -> className + "类代理失败!");
            throw new RuntimeException(e);
        }
    }

    private void doTransform(String key, String className, CtClass cl, Map<String, List<MethodOverWriteConfig>> overWriteConfig) throws CannotCompileException, NotFoundException {
        List<MethodOverWriteConfig> config = overWriteConfig.get(key);
        for (int i = 0; i < config.size(); i++) {
            MethodOverWriteConfig methodOverWriteConfig = config.get(i);
            //1.方法增强
            if (methodOverWriteConfig.getType().equals("method")) {
                doTransformMethod(key, className, cl, methodOverWriteConfig);
                //2.构造方法增强
            } else if (methodOverWriteConfig.getType().equals("constructor")) {
                doTransformConstructor(key, className, cl, methodOverWriteConfig);
            }
        }
    }

    private void doTransformConstructor(String key, String className, CtClass cl, MethodOverWriteConfig methodOverWriteConfig) throws CannotCompileException {
        CtConstructor[] declaredConstructors = cl.getDeclaredConstructors();
        if (declaredConstructors == null || declaredConstructors.length == 0) {
            return;
        }
        //1.单纯的输出
        if (methodOverWriteConfig.getMethodDesc().equals("###")) {
            printConstructorInfo(key, className, declaredConstructors);
            return;
        }
        //2.代码增强
        changeConstructor(declaredConstructors, methodOverWriteConfig);
    }

    private void changeConstructor(CtConstructor[] declaredConstructors, MethodOverWriteConfig methodOverWriteConfig) throws CannotCompileException {
        for (int j = 0; j < declaredConstructors.length; j++) {
            CtConstructor declaredConstructor = declaredConstructors[j];
            if (!methodOverWriteConfig.getMethodDesc().equals(declaredConstructor.getMethodInfo().toString())) {
                continue;
            }
            if (methodOverWriteConfig.getTarget().equals("before")) {
                declaredConstructor.insertBefore(methodOverWriteConfig.getOverWriteText());
            }
            if (methodOverWriteConfig.getTarget().equals("overwrite")) {
                declaredConstructor.setBody(methodOverWriteConfig.getOverWriteText());
            }
            if (methodOverWriteConfig.getTarget().equals("after")) {
                declaredConstructor.insertAfter(methodOverWriteConfig.getOverWriteText());
            }
            break;
        }
    }

    private void printConstructorInfo(String key, String className, CtConstructor[] declaredConstructors) {
        for (int j = 0; j < declaredConstructors.length; j++) {
            CtConstructor declaredConstructor = declaredConstructors[j];
            System.out.println(className + "-" + key + ": 【" + declaredConstructor.getMethodInfo().toString() + "】");
        }
    }

    private void doTransformMethod(String key, String className, CtClass cl, MethodOverWriteConfig methodOverWriteConfig) throws NotFoundException, CannotCompileException {
        CtMethod[] methods = cl.getDeclaredMethods(key);
        if (methods == null || methods.length == 0) {
            return;
        }
        //1.单纯的输出
        if (methodOverWriteConfig.getMethodDesc().equals("###")) {
            printMethodInfo(key, className, methods);
            return;
        }
        //2.进行代码增强
        changeMethod(className, cl, methods, methodOverWriteConfig);
    }

    private void changeMethod(String className, CtClass cl, CtMethod[] methods, MethodOverWriteConfig methodOverWriteConfig) throws CannotCompileException, NotFoundException {
        for (int j = 0; j < methods.length; j++) {
            try {
                CtMethod method = methods[j];
                if (!methodOverWriteConfig.getMethodDesc().equals(method.getMethodInfo().toString())) {
                    continue;
                }
                if (methodOverWriteConfig.getTarget().equals("before")) {
                    method.insertBefore(methodOverWriteConfig.getOverWriteText());
//                    logger.info(String.format("类%s.%s方法代理成功.", className, method.getMethodInfo().toString()));
                    System.out.println(String.format("类%s.%s方法代理成功.", className, method.getMethodInfo().toString()));
                }
                if (methodOverWriteConfig.getTarget().equals("overwrite")) {
                    if (EarlyClassCacheTransformer.getSington().contains(className)) {
                        logger.info("class " + className + " " + method.getMethodInfo().toString() + " 方法在转换器生效前加载，无法复制备份原始方法");
                        method.setBody(methodOverWriteConfig.getOverWriteText());
//                        logger.info(String.format("类%s.%s方法代理成功.", className, method.getMethodInfo().toString()));
                        System.out.println(String.format("类%s.%s方法代理成功.", className, method.getMethodInfo().toString()));
                        method.getMethodInfo().getCodeAttribute().computeMaxStack();
                    } else {
                        if (methodOverWriteConfig.getNeedCopyMethod()) {
                            CtMethod agentMethod = CtNewMethod.copy(method, method.getName() + "$agent", cl, null);
                            agentMethod.setModifiers(Modifier.PUBLIC);
                            cl.addMethod(agentMethod);
                        }

                        method.setBody(methodOverWriteConfig.getOverWriteText());
//                        logger.info(String.format("类%s.%s方法代理成功.", className, method.getMethodInfo().toString()));
                        System.out.println(String.format("类%s.%s方法代理成功.", className, method.getMethodInfo().toString()));
                    }
                }
                if (methodOverWriteConfig.getTarget().equals("after")) {
                    //如果方法没有返回值
                    if ("void".equals(method.getReturnType().getName())) {
                        method.insertAfter(methodOverWriteConfig.getOverWriteText());
//                        logger.info(String.format("类%s.%s方法代理成功.", className, method.getMethodInfo().toString()));
                        System.out.println(String.format("类%s.%s方法代理成功.", className, method.getMethodInfo().toString()));
                    } else {
                        //在有返回值的情况下，暂时无法处理
                        //当然其实我们知道用$1 $2来代表参数，所以其实参数个数有限时，我们是可以根据参数的个数去动态组装调用代码的
                    }
                }
                break;
            } catch (Exception e) {
                e.printStackTrace();
                System.exit(99);
            }
        }
    }

    private void printMethodInfo(String key, String className, CtMethod[] methods) {
        for (int j = 0; j < methods.length; j++) {
            System.out.println(className + "-" + key + ": 【" + methods[j].getMethodInfo().toString() + "】");
        }
    }

    public boolean contains(Class cls) {
        return this.modifyClassConfig.getMethodOverWriteConfigs().get(cls.getName()) != null;
    }
}
